#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
机器人关节数据重放脚本
从日志文件中提取关节数据，并按时间顺序重放
"""

import re
import numpy as np
import time
import argparse
import logging
import sys
from teleop.robot_control.fangnuo_car import FangnuoCar

# 设置日志格式
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler(sys.stdout)
    ]
)

logger = logging.getLogger('JointReplay')

def extract_joint_data(log_file):
    """
    从日志文件中提取关节数据
    返回：包含时间戳和关节数据的列表
    """
    joint_data = []
    timestamp_pattern = r'(\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2},\d{3})'
    joint_pattern = r'更新机器人:\[(.*?)\]'
    
    with open(log_file, 'r') as f:
        content = f.read()
    
    # 查找所有匹配项
    matches = re.finditer(f'{timestamp_pattern}.*?{joint_pattern}', content, re.DOTALL)
    
    for match in matches:
        timestamp_str = match.group(1)
        joint_str = match.group(2)
        
        # 将时间字符串转换为时间戳
        timestamp = time.mktime(time.strptime(timestamp_str, '%Y-%m-%d %H:%M:%S,%f'))
        
        # 将关节数据字符串转换为numpy数组
        joint_values = np.array([float(x) for x in joint_str.split()])
        
        joint_data.append((timestamp, joint_values))
    
    # 按时间戳排序
    joint_data.sort(key=lambda x: x[0])
    
    return joint_data

def replay_joint_data(joint_data, speed_factor=1.0, use_real_arms=True):
    """
    重放关节数据
    参数：
        joint_data: 包含时间戳和关节数据的列表
        speed_factor: 播放速度因子，大于1加速，小于1减速
        use_real_arms: 是否使用真实机械臂
    """
    if not joint_data:
        logger.error("没有找到关节数据，无法重放")
        return
    
    # 初始化机械臂控制器
    logger.info("初始化机械臂控制器...")
    arm_ctrl = FangnuoCar(use_real_arms=use_real_arms)
    
    # 设置机械臂速度逐渐增加到最大值
    logger.info("设置机械臂速度逐渐增加到最大值...")
    arm_ctrl.speed_gradual_max(t=5.0)
    
    # 等待用户确认
    input("机械臂准备就绪，按回车键开始重放...")
    
    # 记录开始时间
    start_time = time.time()
    ref_start_time = joint_data[0][0]
    
    # 重放数据
    logger.info("开始重放关节数据...")
    for i, (timestamp, joint_values) in enumerate(joint_data):
        # 计算应该等待的时间
        if i > 0:
            time_diff = (timestamp - joint_data[i-1][0]) / speed_factor
            time.sleep(max(0, time_diff))
        
        # 控制机械臂
        logger.info(f"重放第 {i+1}/{len(joint_data)} 组数据: {joint_values}")
        arm_ctrl.ctrl_dual_arm(joint_values, np.zeros_like(joint_values))
    
    logger.info("关节数据重放完成")
    
    # 等待用户确认后回到初始位置
    input("重放完成，按回车键让机械臂回到初始位置...")
    arm_ctrl.ctrl_dual_arm_go_home()
    logger.info("机械臂已回到初始位置")

def save_joint_data_to_file(joint_data, output_dir):
    """
    将关节数据保存到文件
    参数：
        joint_data: 包含时间戳和关节数据的列表
        output_dir: 输出目录
    返回：
        csv_file: CSV文件路径
        json_file: JSON文件路径
    """
    import os
    import csv
    import json
    import datetime
    
    # 创建输出目录
    os.makedirs(output_dir, exist_ok=True)
    
    # 生成时间戳作为文件名
    timestamp = datetime.datetime.now().strftime('%Y%m%d_%H%M%S')
    csv_file = os.path.join(output_dir, f'joint_data_{timestamp}.csv')
    json_file = os.path.join(output_dir, f'joint_data_{timestamp}.json')
    
    # 保存为CSV格式
    with open(csv_file, 'w', newline='') as f:
        writer = csv.writer(f)
        # 写入表头
        header = ['timestamp', 'datetime']
        for i in range(len(joint_data[0][1])):
            header.append(f'joint_{i}')
        writer.writerow(header)
        
        # 写入数据
        for timestamp, joint_values in joint_data:
            row = [timestamp, datetime.datetime.fromtimestamp(timestamp).strftime('%Y-%m-%d %H:%M:%S.%f')]
            row.extend(joint_values)
            writer.writerow(row)
    
    # 保存为JSON格式
    json_data = []
    for timestamp, joint_values in joint_data:
        json_data.append({
            'timestamp': timestamp,
            'datetime': datetime.datetime.fromtimestamp(timestamp).strftime('%Y-%m-%d %H:%M:%S.%f'),
            'joint_values': joint_values.tolist()
        })
    
    with open(json_file, 'w') as f:
        json.dump(json_data, f, indent=2)
    
    logger.info(f"关节数据已保存到CSV文件: {csv_file}")
    logger.info(f"关节数据已保存到JSON文件: {json_file}")
    
    return csv_file, json_file

def main():
    parser = argparse.ArgumentParser(description='机器人关节数据重放工具')
    parser.add_argument('--log_file', type=str, default='/home/dw/workspaces/fangnuo/avp_teleoperate/logs/debug-data.log',
                        help='日志文件路径')
    parser.add_argument('--speed', type=float, default=1.0,
                        help='播放速度因子，大于1加速，小于1减速')
    parser.add_argument('--sim', action='store_true',
                        help='使用模拟机械臂而不是真实机械臂')
    parser.add_argument('--save_only', action='store_true',
                        help='仅保存数据，不重放')
    parser.add_argument('--output_dir', type=str, default='./joint_data',
                        help='数据保存目录')
    
    args = parser.parse_args()
    
    # 提取关节数据
    logger.info(f"从日志文件 {args.log_file} 中提取关节数据...")
    joint_data = extract_joint_data(args.log_file)
    logger.info(f"共提取到 {len(joint_data)} 组关节数据")
    
    # 保存关节数据到文件
    csv_file, json_file = save_joint_data_to_file(joint_data, args.output_dir)
    
    # 如果不是仅保存模式，则重放关节数据
    if not args.save_only:
        replay_joint_data(joint_data, args.speed, not args.sim)

if __name__ == '__main__':
    main()
